import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, LSTM
from sklearn.preprocessing import MinMaxScaler

# 1. 生成过去 100 天的销量数据
np.random.seed(42)  # 设置随机种子，保证结果可复现
days = np.arange(100)
# 生成正弦波数据并添加噪声
sales = np.sin(0.1 * days) * 50 + np.random.normal(0, 5, size=100)  # 50为振幅，噪声为标准差5

# 2. 将数据存放在 DataFrame 中
data = pd.DataFrame({'day': days, 'sales': sales})
data.set_index('day', inplace=True)

# 3. 数据预处理
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_data = scaler.fit_transform(data)

# 4. 创建 LSTM 数据集
def create_dataset(data, time_step=1):
    """
    创建用于训练 LSTM 模型的数据集。

    参数:
        data (numpy.ndarray): 输入的时间序列数据。
        time_step (int): 每个样本的时间步数，默认为 1。

    返回:
        X (numpy.ndarray): 输入数据的样本，形状为 (样本数, 时间步数, 特征数)。
        Y (numpy.ndarray): 对应的目标值，形状为 (样本数, 特征数)。

    """
    X, Y = [], []
    for i in range(len(data) - time_step - 1):
        X.append(data[i:(i + time_step), 0])
        Y.append(data[i + time_step, 0])
    return np.array(X), np.array(Y)


time_step = 10
X, Y = create_dataset(scaled_data, time_step)

# reshape input to be [samples, time steps, features]
X = X.reshape(X.shape[0], X.shape[1], 1)

# 5. 创建 LSTM 模型
model = Sequential()
model.add(LSTM(50, return_sequences=True, input_shape=(X.shape[1], 1)))
model.add(LSTM(50, return_sequences=False))
model.add(Dense(1))

# 编译模型
model.compile(optimizer='adam', loss='mean_squared_error')

# 6. 训练模型
model.fit(X, Y, epochs=100, batch_size=1, verbose=2)

# 7. 进行预测
# 先使用训练数据进行预测
train_predict = model.predict(X)

# 8. 反归一化数据
train_predict = scaler.inverse_transform(train_predict)

# 9. 可视化结果
plt.figure(figsize=(10, 6))

# 绘制过去的销量数据
plt.plot(data.index, data['sales'], label='Actual Sales')

# 绘制预测的销量数据
plt.plot(data.index[time_step:len(train_predict) + time_step], train_predict, label='LSTM Predicted Sales', linestyle='--', color='orange')

# 标题和标签
plt.title('Sales Prediction Using LSTM')
plt.xlabel('Day')
plt.ylabel('Sales')
plt.legend()
plt.grid(True)

# 保存图像到文件
plt.savefig('lstm_sales_prediction.png')

# 显示图像
plt.show()
